Description

When DelayedScaling is used with multiple Transformer Engine modules under activation checkpointing, FP8 global state update ownership is incorrectly handled during recomputation.

Checkpoint recomputation introduces nested FP8 autocast scopes. These scopes can independently advance forward amax/scale state, causing one forward update per recomputed checkpoint segment instead of one update per logical forward pass.

With reentrant checkpointing, recomputed modules can also repeatedly reacquire first-module ownership. This causes backward amax reduction and scale updates to execute once per recomputed module instead of once per logical backward pass.

As a result, FP8 state may be advanced multiple times within a single logical microbatch.

Minimal reproduction

The issue reproduces using only public PyTorch and Transformer Engine APIs:

• transformer_engine.pytorch.Linear
• Transformer Engine checkpointing
• FP8 autocast
• DelayedScaling
• Format.HYBRID
• Three sequential te.Linear modules

The reproduction does not depend on a model framework, distributed runtime, external dataset, saved checkpoint, or application-specific code.

On unmodified Transformer Engine main at commit 70957ad5aa8337a12c254e86a6e0b7ffdbbd4b59:

Expected for every configuration:

forward FP8 state updates:  1
backward FP8 state updates: 1

The loss, input gradient, and all weight gradients remain finite and nonzero in the minimal reproduction. The incorrect state progression therefore does not necessarily produce an immediate runtime error.

The results were reproduced across three seeds before applying the fix.

Root cause

is_first_fp8_module is stored in process-global FP8 state, while ownership of the corresponding update belongs to the surrounding logical FP8 autocast and checkpoint frame.

During activation recomputation:

1. Each recomputation autocast scope can be treated as a new forward update boundary.
2. Reentrant checkpoint frames restore saved first-module state.
3. First-module ownership can be restored after it has already been queried and consumed.
4. Subsequent recomputed modules can therefore acquire the same ownership again.

This causes:

• forward FP8 state updates once per recomputed checkpoint segment;
• backward FP8 state reduction and update once per recomputed module under reentrant checkpointing.

The affected state transitions involve:

transformer_engine/pytorch/distributed.py
transformer_engine/pytorch/quantization.py
transformer_engine/pytorch/module/linear.py
transformer_engine/pytorch/module/layernorm_linear.py
transformer_engine/pytorch/module/layernorm_mlp.py

Fix

The fix makes FP8 update ownership checkpoint-frame-aware and preserves it across activation recomputation.

The implementation:

• saves and restores first-module ownership independently for each checkpoint frame;
• reserves the single backward-update ownership during the original no-grad forward of a reentrant checkpoint;
• allows recomputation to consume that ownership normally without restoring it after every module;
• prevents recomputation autocast scopes from establishing additional logical forward-update boundaries;
• prevents recomputation autocast exit from independently advancing forward DelayedScaling state;
• restores the surrounding FP8 autocast and recomputation state when a checkpoint frame exits;
• applies the checkpoint-aware ownership handling consistently to Linear, LayerNormLinear, and LayerNormMLP.

After the change, Transformer Engine performs exactly one forward and one backward global FP8 state update for each logical microbatch.

No manual FP8 update call or application-level workaround is required.

Tests

A regression test is added at:

tests/pytorch/test_reentrant_fp8_updates.py

The tests cover:

• no checkpoint;
• non-reentrant checkpointing with one segment;
• non-reentrant checkpointing with one segment per layer;
• reentrant checkpointing with one segment;
• reentrant checkpointing with one segment per layer;
• nested reentrant checkpointing;
• restoration of surrounding FP8 autocast state;
• finite and nonzero input and weight gradients.

Before the fix:

6 failed, 2 passed

All six failures were caused by forward or backward update counts differing from 1.

After the fix:

8 passed

The complete six-path matrix was also run with five seeds:

30/30 runs passed
forward updates:  1
backward updates: 1
loss: finite
input gradients: finite and nonzero
weight gradients: finite and nonzero
surrounding FP8 state: restored

Related existing tests also passed:

• full activation recompute with reentrant and non-reentrant DelayedScaling;
• reentrant activation recompute with CurrentScaling;
• recipe, scale, and quantizer-state tests;
• Linear, LayerNormLinear, and LayerNormMLP backward tests;
• nested FP8 autocast and torch.compile coverage;
• PyTorch L0 lint;
• license L0;
• Black, syntax, and whitespace checks.

The full tests/pytorch/test_checkpoint.py compatibility suite was not runnable locally because the checkout did not contain its required pre-generated .pt artifacts. The affected cases exited with FileNotFoundError before entering the Transformer Engine checkpoint-loading path.

Real-world impact

The issue was discovered during single-GPU pretraining of a 1.68B-parameter decoder model on a 24 GB NVIDIA GeForce RTX 5090 D v2. Reentrant activation checkpointing was required to fit the workload in device memory.

The model contained 196 Transformer Engine Linear modules. Before the workaround, each microbatch executed:

forward FP8 state updates:  29
backward FP8 state updates: 196

With gradient accumulation over 64 microbatches, this resulted in:

forward updates per optimizer step:   1,856
backward updates per optimizer step: 12,544

Because FP8 amax and scale state was advanced while backward propagation was still in progress, later modules in the same backward pass could consume different state from earlier modules.

This produced execution-order-sensitive dgrad corruption. Observed examples included transitions from gradients on the order of 1e-6 to values above 300, and intermittent pre-clipping global gradient-norm spikes as high as 629406.

The affected runs continued to report finite losses, making the failure easy to miss when only loss values or post-clipping gradients were monitored.

After restoring one forward and one backward update per logical microbatch:

• update counts became exactly 1/1;
• the previously observed gradient explosions no longer occurred;
• measured gradient norms were within approximately 2% of the BF16 reference in the tested steps;
• parameter update norms remained stable;
• the FP8 throughput benefit was preserved at approximately 13.2k tokens/s.

This workload is included only as additional impact validation. The regression tests and minimal reproduction use public Transformer Engine APIs exclusively.

Diagnostic behavior

The standard Transformer Engine debug instrumentation did not reproduce the numerical explosion in this workload.

However, enabling it reduced throughput from approximately 13.0k tokens/s to 3.2k tokens/s. It also changed quantizer execution properties and substantially altered launch timing.

The debug run therefore could not exclude an execution-order-sensitive state-management issue.

Duplicate state updates were instead confirmed using lightweight Python-side counting of the actual FP8 update calls. The counter did not:

• collect tensor values;
• install tensor backward hooks;
• invoke CUDA synchronization;
• enable debug fallback paths.

Environment

Initial discovery baseline: Transformer Engine 2.16.0
Current main validation:    2.18.0.dev0+70957ad5
PyTorch:                    2.13.0+cu130
CUDA toolkit:               13.0.88
cuDNN:                      9.20.0.48
cuBLAS:                     13.1.1.3
GPU:                        NVIDIA GeForce RTX 5090 D v2, 24 GB
Compute capability:         12.0 (SM120)
Execution mode:             Single GPU

The bookkeeping failure occurs in Transformer Engine framework state management. It has been reproduced on compute capability 12.0, but there is currently no evidence that it is specific to SM120.

Type of change

• Documentation change
• Bug fix
• New feature
• Breaking change
• Infra/Build change
• Code refactoring

Checklist

• I have read and followed the contributing guidelines.
• The functionality is complete.
• I have commented the hard-to-understand state-management logic.
• I have made corresponding documentation changes — not applicable; no public API or user-facing behavior was added.
• My changes generate no new warnings.
• I have added tests that prove the fix is effective.
• New and relevant existing unit tests pass locally with my changes.